This PR adds parallel folding for AutoEP: tensor parallelism (AutoTP) for the dense/attention path can now coexist with expert parallelism (AutoEP) for the routed-expert path on the same set of ranks, without forcing EP to be a subset of DP — an EP group may span TP lanes and dense-DP ranks (cross-lane EP).
(This PR should be adjusted for ZeRO3 support after #8060 is merged)

Design

Attention/dense and MoE are treated as two independent partitionings of the same rank set, parameterized per parameter family:

• Dense / attention / shared-expert params: stage_size = tp * dp
• Routed-expert params: stage_size = ep * etp * edp

dp and edp are always derived, never user-configured, so the invariant tp * dp == ep * etp * edp == stage_size cannot be broken from config. The only structural requirement is that the expert width tiles the stage (stage_size % (ep * etp) == 0); EP groups are then laid across a TP-lane-major rank ordering, so they may span TP lanes and dense-DP ranks.

Configuration

No new config section. Folding is expressed by the coexistence of the existing tensor_parallel and expert_parallel sections:

{
  "tensor_parallel":  { "autotp_size": 4 },
  "expert_parallel":  { "enabled": true, "autoep_size": 4,
                        "expert_tensor_parallel_size": 1 }
}

expert_tensor_parallel_size is carried as a config field but currently must be 1 (expert-internal TP is reserved as follow-up and rejected fail-fast). Validation enforces stage divisibility, TP/sequence-parallel exclusivity, and preset_model consistency between the two sections.

Cross-lane expert parallelism

Expert parallelism no longer has to be a subset of data parallelism. Shapes where the expert width exceeds (or does not divide) the dense data-parallel size are supported, for example:

• world=4, TP=4, EP=4 (dp=1): the EP group is the whole TP group — one expert per rank.
• world=4, TP=2, EP=4 (dp=2): the EP group spans both TP lanes and both DP ranks.
• world=8, TP=4, EP=4 (dp=2, edp=2): EP groups span TP lanes with expert replication.

The per-family gradient convention is keyed to each parameter's replication structure, not to the EP layout, so it holds across the whole tp*dp pool:

• Router/gate and dense/LayerNorm are AVERAGE over the TP (token-replication) group. The folded router runs redundantly on every TP peer; its partitioned work is reconstructed into a replicated full view by restore_combined, whose all-gather backward injects a tp_size factor that AVERAGE divides out.
• Routed experts cancel that same restore_combined tp_size factor (divide by tp_size, no TP all-reduce) and reduce data-parallel over the expert-data-parallel (EDP) group. Without the cancellation, folded expert gradients are over-scaled by tp_size — invisible to scale-invariant Adam, but real for non-adaptive optimizers and for gradient clipping (it inflates the expert contribution to the global grad norm). This is now fixed for all folded shapes (the MVP TP2×EP4 shape included).

What's included

• Folded process-group derivation using the generalized expert/data-parallel group creation (mp_mode TP-strided vs SP-consecutive ordering), including cross-lane EP group tables.
• Route-full / partition-dispatch path for folded MoE (deepspeed/moe/ep_tp_dispatch.py), with AutoTP skipping AutoEP subtrees.
• Per-family folded gradient reduction: AVERAGE for replicated router/gate and dense/LayerNorm; a dedicated tp_size cancellation for routed experts; SKIP for genuinely TP-sharded params; SUM contracts reserved for a future true sequence-parallel path.
• Per-parameter-family ZeRO checkpoint metadata (routed-expert vs dense/router/shared placement) and folded ZeRO-1/2 optimizer-state handling.

Correctness & validation

• Router/gate and LayerNorm gradient parity against a non-folded ZeRO baseline (atol=1e-1, rtol=5e-3, fp32), on TP2×EP4 (8-rank) and the cross-lane shapes TP2×EP4 (4-rank, ep>dp) and TP4×EP4 (4-rank, dp=1); scale 1.0.
• Routed-expert weight parity against a non-folded baseline, verified with SGD (Adam is scale-invariant and would mask a uniform gradient-scale error), for the MVP TP2×EP4 shape and cross-lane TP4×EP4 with edp=1 and edp=2.
• New folding unit tests for config, cross-lane group layout, dispatch, runtime, gradient parity, and checkpoint save/load (multi-rank GPU cases gated for GPU runners; CPU/Gloo parity runs on CI).
• Real-H100 confirmation (8×H100): router/gate, LayerNorm, and routed-expert gradient parity to the non-folded baseline hold (scale 1.0) for MVP TP2×EP4 and cross-lane TP4×EP4 (edp=1 and edp=2); cross-lane folded training runs with finite loss and finite, non-zero expert/router gradients.
• Passes the full unit test suite (aws-torch-latest-full) on H100 GPUs.

Scope / follow-ups

• This PR covers AutoEP + AutoTP folding, including cross-lane EP with etp=1. The replicated-grad reduction is mode-aware so the sequence-parallel (Ulysses) folding case fits the same contract; AutoTP + AutoEP is the validated path here.
• Expert-internal tensor parallelism (expert_tensor_parallel_size > 1) is reserved for a follow-up.
• ZeRO-3 composition with folding is planned as separate follow-up work (after Support AutoEP with ZeRO-3 zero.Init source modules #8060 is merged).